#include "opencv2/opencv.hpp"
#include "boost/filesystem.hpp"

using namespace std;
using namespace cv;

#define IS_RED(x) (x[0] == 0 && x[1] == 0 && x[2] == 255)
#define IS_BLUR(x) (x[0] == 255 && x[1] == 0 && x[2] == 0)

#define IS_BEACON(img, x, y) ( \
	IS_RED(img.ptr<Vec3b>(y-1)[x-1]) && IS_RED(img.ptr<Vec3b>(y-1)[x]) && IS_RED(img.ptr<Vec3b>(y-1)[x+1]) && \
	IS_RED(img.ptr<Vec3b>(y)[x-1]) && IS_RED(img.ptr<Vec3b>(y)[x]) && IS_RED(img.ptr<Vec3b>(y)[x+1]) && \
	IS_RED(img.ptr<Vec3b>(y+1)[x-1]) && IS_RED(img.ptr<Vec3b>(y+1)[x]) && IS_RED(img.ptr<Vec3b>(y+1)[x+1]))

#define IS_VALID_POS(img, x, y) ( \
	IS_BLUR(img.ptr<Vec3b>(y-1)[x-1]) && IS_BLUR(img.ptr<Vec3b>(y-1)[x]) && IS_BLUR(img.ptr<Vec3b>(y-1)[x+1]) && \
	IS_BLUR(img.ptr<Vec3b>(y)[x-1]) && IS_BLUR(img.ptr<Vec3b>(y)[x]) && IS_BLUR(img.ptr<Vec3b>(y)[x+1]) && \
	IS_BLUR(img.ptr<Vec3b>(y+1)[x-1]) && IS_BLUR(img.ptr<Vec3b>(y+1)[x]) && IS_BLUR(img.ptr<Vec3b>(y+1)[x+1]))

enum LANDMARK_TYPE
{
	BEACON,
	VALID_POS,
};

struct Landmark
{
	LANDMARK_TYPE type;
	int img_x, img_y;
	float map_x, map_y;
	int idx;
};

void loadSize(std::string size_txt, float &map_w, float &map_h);
std::vector<Landmark> findLandmarks(const cv::Mat &config_img, int start_idx, float map_w, float map_h);
cv::Mat drawLandmarks(const cv::Mat &config_img, const std::vector<Landmark> &landmarks);
void writeLandmarksTxt(std::string beacon_txt, const std::vector<Landmark> &landmarks, LANDMARK_TYPE type);

int main(int argc, char **argv)
{
	// handle input
	string work_dir = argc > 1 ? argv[1] : "";
	boost::filesystem::path p_root(work_dir);
	
	string config_png = argc > 2 ? argv[2] : "";
	config_png = (p_root / config_png).string();
	boost::filesystem::path p_config(config_png);

	int start_idx = argc > 3 ? atoi(argv[3]) : 1;

	string beacon_txt = (p_root / p_config.filename().replace_extension(".beacon.txt")).string();
	string beacon_png = (p_root / p_config.filename().replace_extension(".beacon.png")).string();
	string valid_pos_txt = (p_root / p_config.filename().replace_extension(".valid_pos.txt")).string();

	if (config_png.empty() || beacon_txt.empty()) {
		printf("usage: [this.exe] [config png] [config txt]\n");
		system("pause"); exit(1);
	}

	// load size
	string size_txt = (p_root / p_config.filename().replace_extension(".size.txt")).string();
	float map_w, map_h;
	loadSize(size_txt, map_w, map_h);

	// load png
	Mat config_img = imread(config_png);
	vector<Landmark> landmarks = findLandmarks(config_img, start_idx, map_w, map_h);
	Mat canvas;
	canvas = drawLandmarks(config_img, landmarks);
	imwrite(beacon_png, canvas);
	writeLandmarksTxt(beacon_txt, landmarks, BEACON);
	writeLandmarksTxt(valid_pos_txt, landmarks, VALID_POS);
}

void loadSize(std::string size_txt, float &map_w, float &map_h)
{
	ifstream ifs(size_txt);
	ifs >> map_w >> map_h;
	ifs.close();
}

std::vector<Landmark> findLandmarks(const cv::Mat &config_img, int start_idx, float map_w, float map_h)
{
	vector<Landmark> landmarks;

	// find beacon
	for (int y = 1; y < config_img.rows - 1; ++y) {
		for (int x = 1; x < config_img.cols - 1; ++x) {
			if (IS_BEACON(config_img, x, y)) {
				Landmark mark;
				mark.type = BEACON;
				mark.img_x = x;
				mark.img_y = y;
				mark.map_x = float(x) / config_img.cols * map_w;
				mark.map_y = float(y) / config_img.rows * map_w;
				mark.idx = landmarks.size() + start_idx;

				landmarks.push_back(mark);
			}
		}
	}

	// find valid pos
	for (int y = 1; y < config_img.rows - 1; ++y) {
		for (int x = 1; x < config_img.cols - 1; ++x) {
			if (IS_VALID_POS(config_img, x, y)) {
				Landmark mark;
				mark.type = VALID_POS;
				mark.img_x = x;
				mark.img_y = y;
				mark.map_x = float(x) / config_img.cols * map_w;
				mark.map_y = float(y) / config_img.rows * map_w;
				mark.idx = -1;

				landmarks.push_back(mark);
			}
		}
	}

	return landmarks;
}

cv::Mat drawLandmarks(const cv::Mat &config_img, const std::vector<Landmark> &landmarks)
{
	Mat canvas = config_img.clone();

	for (int i = 0; i < landmarks.size(); ++i) {
		const Landmark &mark = landmarks[i];

		// beacon
		if (mark.type == BEACON) {
			circle(canvas, Point(mark.img_x, mark.img_y), 10, CV_RGB(255, 0, 0), 2);
			putText(canvas, to_string(mark.idx), Point(mark.img_x, mark.img_y - 15), 0, 1, CV_RGB(255, 0, 0), 1);
		} 

		// valid pos
		if (mark.type == VALID_POS) {
			circle(canvas, Point(mark.img_x, mark.img_y), 8, CV_RGB(0, 0, 255), 2);
		}
	}

	return canvas;
}

void writeLandmarksTxt(std::string beacon_txt, const std::vector<Landmark> &landmarks, LANDMARK_TYPE type)
{
	ofstream ofs(beacon_txt);
	if (!ofs.is_open())
		return;

	for (int i = 0; i < landmarks.size(); ++i) {
		const Landmark &mark = landmarks[i];
		if (mark.type == type) {
			ofs << mark.idx 
				<< ", " << mark.img_x << ", " << mark.img_y
				<< ", " << mark.map_x << ", " << mark.map_y
				<< endl;
		}
	}

	ofs.close();
}
